import pymongo
from bson.objectid import ObjectId

class UserManager:
    def __init__(self, config):
        self.config = config
        # 修改这里，使用正确的配置路径
        # 检查配置格式并提供兼容性
        if 'mongodb_uri' in config:
            mongodb_uri = config['mongodb_uri']
        elif 'mongodb' in config.get('database', {}) and 'connection_string' in config.get('database', {}).get('mongodb', {}):
            mongodb_uri = config['database']['mongodb']['connection_string']
        else:
            # 使用默认连接字符串
            mongodb_uri = "mongodb://localhost:27017/"
            
        self.client = pymongo.MongoClient(mongodb_uri)
        # 同样检查数据库名称配置
        if 'database' in config:
            db_name = config.get('database', 'stock_ai_analysis')
        else:
            db_name = 'stock_ai_analysis'
            
        self.db = self.client[db_name]
        self.users_collection = self.db[config.get('db_collection', 'users')]
        
    def create_user(self, user_data):
        """创建新用户"""
        required_fields = ['username', 'email']
        for field in required_fields:
            if field not in user_data:
                raise ValueError(f"Missing required field: {field}")
                
        # 检查用户名是否已存在
        if self.users_collection.find_one({'username': user_data['username']}):
            raise ValueError("Username already exists")
            
        # 设置默认值
        user_data.setdefault('notification_methods', ['email'])
        user_data.setdefault('strategy_preference', ['short', 'medium'])
        user_data.setdefault('subscribed_stocks', [])
        
        return self.users_collection.insert_one(user_data).inserted_id
        
    def get_user(self, user_id):
        """获取用户信息"""
        return self.users_collection.find_one({'_id': ObjectId(user_id)})
        
    def update_user(self, user_id, update_data):
        """更新用户信息"""
        return self.users_collection.update_one(
            {'_id': ObjectId(user_id)},
            {'$set': update_data}
        )
        
    def subscribe_stock(self, user_id, stock_code):
        """订阅股票"""
        return self.users_collection.update_one(
            {'_id': ObjectId(user_id)},
            {'$addToSet': {'subscribed_stocks': stock_code}}
        )
        
    def unsubscribe_stock(self, user_id, stock_code):
        """取消订阅股票"""
        return self.users_collection.update_one(
            {'_id': ObjectId(user_id)},
            {'$pull': {'subscribed_stocks': stock_code}}
        )
        
    def get_stock_subscribers(self, stock_code):
        """获取订阅特定股票的所有用户"""
        return list(self.users_collection.find({'subscribed_stocks': stock_code}))
        
    def set_notification_preferences(self, user_id, methods):
        """设置通知方式偏好"""
        return self.users_collection.update_one(
            {'_id': ObjectId(user_id)},
            {'$set': {'notification_methods': methods}}
        )
        
    def set_strategy_preferences(self, user_id, preferences):
        """设置策略偏好"""
        return self.users_collection.update_one(
            {'_id': ObjectId(user_id)},
            {'$set': {'strategy_preference': preferences}}
        ) 